#include <iostream>
#include <iomanip>
#include <sstream>

#include "wad.hpp"
#include "tool.hpp"

using namespace std;

wad_t::wad_t()
{
	this->init();
}

wad_t::wad_t(string wadfile)
{
	this->init();
	this->open(wadfile);
}

wad_t::~wad_t()
{
	this->close();
}

void wad_t::init()
{
	opened = false;

	wadtype = 0;

	content = NULL;
	for(size_t i = 0; i < sizeof(titlekey); i++)
		titlekey[i] = 0x00;
}

bool wad_t::open(string wadfile)
{
	if(this->is_open())
		this->close();

	wadstream.open(wadfile, ios_base::in | ios_base::binary);
	if(wadstream.fail())
	{
		cerr << "File couldn't be found: " << wadfile << endl;
		this->init();
		wadstream.close();
		return false;
	}

	header.open(wadstream, 0x00, 0x00);
	if(header.len != 0x20)
	{
		cerr << "Invalid WAD (" << wadfile << "): ";
		cerr << "header size is 0x" << hex << header.len << endl;
		this->init();
		wadstream.close();
		return false;
	}

	wadtype = get32(wadstream);
	if(wadtype != 0x49730000 && wadtype != 0x69620000)
	{
		cerr << "Invalid WAD (" << wadfile << "): ";
		cerr << "unknown type " << hex << wadtype << endl;
		this->init();
		wadstream.close();
		return false;
	}

	opened = true;

	cert.open(wadstream, 0x08, header.pos + add_padding(header.len, 0x40));
	tik.open(wadstream, 0x10, cert.pos + add_padding(cert.len, 0x40));
	tmd.open(wadstream, 0x14, tik.pos + add_padding(tik.len, 0x40));
	contents.open(wadstream, 0x18, tmd.pos + add_padding(tmd.len, 0x40));
	footer.open(wadstream, 0x1c, contents.pos + add_padding(contents.len, 0x40));

	//wadstream.seekg(tikpos + 0x01bf);
	//wadstream.read(titlekey, 16);
	this->decrypt_title_key(titlekey);

	wadstream.seekg(tmd.pos + 0x01de);
	contentcount = get16(wadstream);
	
	try
	{
		content = new content_t[contentcount];
	}
	catch(exception& e)
	{
		cerr << "Error: " << e.what() << endl;
		this->close();
		return false;
	}
	uint32_t pos = contents.pos;
	for(uint16_t i=0; i < contentcount; i++)
	{
		content[i].open(wadstream, tmd.pos + 0x24*i, pos);
		pos += add_padding(content[i].len, 0x40);
	}
	return true;
}

void wad_t::close()
{
	opened = false;
	wadstream.close();
	delete[] content;
}

bool wad_t::is_open()
{
	return opened;
}

void wad_t::decrypt_title_key(uint8_t* titlekey)
{
	uint8_t encrypted_key[16];

	uint8_t common_key[16] = {0xeb,0xe4,0x2a,0x22,
	                          0x5e,0x85,0x93,0xe4,
	                          0x48,0xd9,0xc5,0x45,
	                          0x73,0x81,0xaa,0xf7};

	uint8_t iv[16] =         {0x00,0x00,0x00,0x00,
	                          0x00,0x00,0x00,0x00,
	                          0x00,0x00,0x00,0x00,
	                          0x00,0x00,0x00,0x00};

	wadstream.seekg(tik.pos + 0x01dc);
	wadstream.read(reinterpret_cast<char*>(iv), 8);

	wadstream.seekg(tik.pos + 0x01bf);
	wadstream.read(reinterpret_cast<char*>(encrypted_key),sizeof(encrypted_key));

	aes_cbc_dec(common_key, iv, encrypted_key, sizeof(encrypted_key), titlekey);
}

bool wad_t::unpack(string destinationpath)
{
	if(!this->is_open())
		return false;

	if (destinationpath[destinationpath.length()-1] != '/')
	{
		destinationpath += "/";
	}

	cout << "Unpacking cert" << endl;
	cert.unpack(wadstream, destinationpath + "/cert.bin");
	cout << "Unpacking tik" << endl;
	tik.unpack(wadstream, destinationpath + "/tik.bin");
	cout << "Unpacking tmd" << endl;
	tmd.unpack(wadstream, destinationpath + "/tmd.bin");
	cout << "Unpacking footer" << endl;
	footer.unpack(wadstream, destinationpath + "/footer.bin");

	for(uint16_t i=0; i<contentcount; i++)
	{
		cout << "Unpacking content " << i << endl;
		ostringstream name;
		name << destinationpath << setw(8) << setfill('0') << hex << content[i].cid;
		content[i].unpack(wadstream, name.str());
		name << ".app";
		cout << "Decrypting content " << i << endl;
		content[i].unpack(wadstream, name.str(), titlekey);
	}
	
	return true; //TODO
}

void wad_t::print_header()
{
	if(!this->is_open())
	{
		cout << "No WAD opened" << endl;
		return;
	}
	cout << "header.pos: " << hex << header.pos << endl;
	cout << "header.len:" << hex << header.len << endl;
	cout << "cert.pos:" << hex << cert.pos << endl;
	cout << "cert.len:" << hex << cert.len << endl;
	cout << "tik.pos:" << hex << tik.pos << endl;
	cout << "tik.len:" << hex << tik.len << endl;
	cout << "titlekey:";
	for(size_t j=0; j<sizeof(titlekey); j++)
		cout << setw(2) << setfill('0') << hex << static_cast<short>(titlekey[j]);
	cout << endl;
	cout << "tmd.pos:" << hex << tmd.pos << endl;
	cout << "tmd.len:" << hex << tmd.len << endl;
	cout << "contents.pos:" << hex << contents.pos << endl;
	cout << "contents.len:" << hex << contents.len << endl;
	cout << "contentcount:" << hex << contentcount << endl;
	cout << endl;

	for(uint16_t i=0; i<contentcount ; i++)
	{
		cout << "content:" << i << endl;
		cout << "pos:" << hex << content[i].pos << endl;
		cout << "cid:" << hex << content[i].cid << endl;
		cout << "index:" << hex << content[i].index << endl;
		cout << "type:" << hex << content[i].type << endl;
		cout << "len:" << hex << content[i].len << endl;
		cout << "hash:";
		for(size_t j=0; j<sizeof(content[i].hash); j++)
			cout << setw(2) << setfill('0') << hex << static_cast<short>(content[i].hash[j]);
		cout << endl << endl;
	}

	cout << "footer.pos:" << hex << footer.pos << endl;
	cout << "footer.len:" << hex << footer.len << endl;
	cout << endl;
}
